#include <mm/slab.h>
#include <mm/mm.h>
#include <lib/math.h>
#include <string.h>

#define LOG_TAG "SLAB"
#include <log.h>

kmem_cache_t main_kmem_cache;

lnode_t kmem_cache_list;

static void init_cache(kmem_cache_t* cache, char* name, size_t obj_size, size_t page_per_cache, int opt) {
    if (!cache) {
        return;
    }
    size_t offset = sizeof(long);

    if (opt & ALIGN_DMA) {
        offset = 128;
    }

    obj_size = ROUNDUP(obj_size, offset);

    *cache = (kmem_cache_t){
        .obj_size = obj_size,
        .page_per_cache = page_per_cache,
        .slab_count = 1,
        .obj_per_slab = (page_per_cache << PG_SIZE_BITS) / (obj_size + sizeof(obj_idx_t)),
    };

    size_t free_list_size = cache->obj_per_slab * sizeof(obj_idx_t);

    cache->offset = ROUNDUP(free_list_size + sizeof(slab_t), offset);

    size_t slab_size = cache->offset + cache->obj_per_slab * obj_size;
    size_t pages_size = page_per_cache << PG_SIZE_BITS;
    size_t overflow_obj_count = 0;
    if (slab_size > pages_size) {
        overflow_obj_count = ICEIL(slab_size - pages_size, obj_size);
    }
    cache->obj_per_slab -= overflow_obj_count;

    strcpy(cache->cache_name, name);

    list_head_init(&cache->free);
    list_head_init(&cache->full);
    list_head_init(&cache->used);
    list_append(&kmem_cache_list, &cache->caches);
}

kmem_cache_t* new_kmem_cache(char* name, size_t obj_size, size_t page_per_cache, int opt) {
    kmem_cache_t* cache = (kmem_cache_t*)take_obj(&main_kmem_cache);

    init_cache(cache, name, obj_size, page_per_cache, opt);

    return cache;
}

static slab_t* alloc_slab(size_t page_per_cache) {
    return (slab_t*)mapping_pages(get_pages(page_per_cache, PG_KERNEL), page_per_cache, PG_WRITE, VMAP_NULL);
}

static slab_t* new_slab(kmem_cache_t* cache) {
    slab_t* slab = alloc_slab(cache->page_per_cache);

    if (!slab) {
        LOGE("kmem_cache: alloc slab failed\n");
        return NULL;
    }

    slab->next_free = slab->used_objs = 0;
    slab->first_obj = (void*)((uptr)slab + cache->offset);

    obj_idx_t* list = slab->free_list;
    for (long i = 0; i < cache->obj_per_slab - 1; i++) {
        list[i] = i + 1;
    }
    list[cache->obj_per_slab - 1] = -1;

    list_append(&cache->free, &slab->slabs);
    return slab;
}

void* take_obj(kmem_cache_t* cache) {
    slab_t* pos;
    if (!list_empty(&cache->used)) {
        pos = list_entry(cache->used.next, typeof(*pos), slabs);
    } else if (list_empty(&cache->free)) {
        pos = new_slab(cache);
    } else {
        pos = list_entry(cache->free.next, typeof(*pos), slabs);
    }

    if (!pos) {
        LOGE("kmem_cache: no obj available\n");
        return NULL;
    }

    obj_idx_t nxt_idx = pos->next_free;
    pos->next_free = pos->free_list[nxt_idx];
    pos->used_objs++;
    cache->used_count++;

    list_remove(&pos->slabs);
    if (pos->next_free == -1) {
        list_append(&cache->full, &pos->slabs);
    } else {
        list_append(&cache->used, &pos->slabs);
    }

    return (void*)((uptr)pos->first_obj + nxt_idx * cache->obj_size);
}

int release_obj(kmem_cache_t* cache, void* obj) {
    slab_t* pos, * nxt;
    obj_idx_t idx;

    lnode_t* lists[2] = { &cache->full, &cache->used };
    for (int i = 0; i < 2; i++) {
        list_for_each(pos, nxt, lists[i], slabs) {
            if (pos->first_obj > obj) {
                continue;
            }
            idx = (obj - pos->first_obj) / cache->obj_size;
            if (idx < cache->obj_per_slab) {
                goto found;
            }
        }
    }

    return 0;

found:
    pos->free_list[idx] = pos->next_free;
    pos->next_free = idx;
    pos->used_objs--;
    cache->used_count--;

    list_remove(&pos->slabs);
    if (pos->used_objs) {
        list_append(&cache->used, &pos->slabs);
    } else {
        list_append(&cache->free, &pos->slabs);
    }
    return 1;
}

void kmem_cache_init() {
    list_head_init(&kmem_cache_list);
    init_cache(&main_kmem_cache, "main-kmem_cache", sizeof(kmem_cache_t), 1, 0);
}

void kmem_cache_show_status() {
    LOGW("\t<name>\t<slab>  <pg/s>  <obj/s>  <used>\n");

    kmem_cache_t* pos, * n;
    list_for_each(pos, n, &kmem_cache_list, caches) {
        LOGI("%-16s %4d %4d %4d %4d\n",
            pos->cache_name,
            pos->slab_count,
            pos->page_per_cache,
            pos->obj_per_slab,
            pos->used_count);
    }
}